#人工智能第三章进阶作业
#1. 在测试视频(OpenCV安装目录\sources\samples\data)上，使用基于混合高斯模型的背景提取算法，提取前景并显示(显示二值化图像，前景为白色)。

import cv2
# 加载视频
cap = cv2.VideoCapture()
cap.open(r'E:\deep learning\cv\data\vtest.avi')
if not cap.isOpened():
    print("无法打开视频文件")
pBgModel = cv2.createBackgroundSubtractorMOG2()

while True:
    flag, source = cap.read()
    if not flag:
        break
    image = cv2.pyrDown(source)    #降低分辨率

    fgMask = pBgModel.apply(image)  #"apply"函数会实时的将"pBgModel"中的当前帧缩放图像image采用pBgModel的方式进行处理，并提取fgMask，即前景的掩膜

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))    # 构造一个全是1的kernel用于形态学的操作
    morphImage_open = cv2.morphologyEx(fgMask, cv2.MORPH_OPEN, kernel, iterations=5)  # 对图像进行开运算
    mask = fgMask - morphImage_open
    _, Mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    backGround = pBgModel.getBackgroundImage()  #背景图像
    cv2.imshow('source', image)
    cv2.imshow('background', backGround)
    cv2.imshow('foreground', Mask)
    key = cv2.waitKey(10)
    if key == 27:
        break

cap.release()
cv2.destroyAllWindows()

##########################################################################################################

#2. 在1基础上，将前景目标进行分割，进一步使用不同颜色矩形框标记，并在命令行窗口中输出每个矩形框的位置和大小。

import cv2
# 加载视频
cap = cv2.VideoCapture()
cap.open(r'E:\deep learning\cv\data\vtest.avi')
if not cap.isOpened():
    print("无法打开视频文件")
pBgModel = cv2.createBackgroundSubtractorMOG2()

def labelTargets(img, mask, threshold):
    seg = mask.copy()
    cnts = cv2.findContours(seg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    count = 0
    for i in cnts[1]:
        area = cv2.contourArea(i)
        if area < threshold:
            continue
        count += 1
        rect = cv2.boundingRect(i)   #返回矩形坐标
        print("矩形：X:{} Y:{} 宽：{} 高：{}".format(rect[0], rect[1], rect[2], rect[3]))
        cv2.drawContours(img, [i], -1, (255, 255, 0), 1)      #画轮廓
        cv2.rectangle(img, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (0, 0, 255), 1)       #画矩形框
        cv2.putText(img, str(count), (rect[0], rect[1]), cv2.FONT_HERSHEY_PLAIN, 0.5, (0, 255, 0))           #标数字
    return count

while True:
    flag, source = cap.read()
    if not flag:
        break
    image = cv2.pyrDown(source)

    fgMask = pBgModel.apply(image)

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    morphImage_open = cv2.morphologyEx(fgMask, cv2.MORPH_OPEN, kernel, iterations=5)
    mask = fgMask - morphImage_open
    _, Mask = cv2.threshold(mask, 30, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # Mask = cv2.GaussianBlur(Mask, (5, 5), 0)
    targets = labelTargets(image, Mask, 30)

    print("共检测%s个目标" % targets)
    backGround = pBgModel.getBackgroundImage()
    foreGround = image - backGround
    cv2.imshow('source', image)
    cv2.imshow('background', backGround)
    cv2.imshow('foreground', Mask)
    key = cv2.waitKey(10)
    if key == 27:
        break

##########################################################################################################

#3. 安装ImageWatch，并在代码中通过设置断点，观察处理中间结果图像。

#未使用vs

#4. 使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。

import cv2

videoFileName = r'E:\deep learning\cv\data\vtest.avi'
# 角点检测参数设定
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)
                      # maxCorners : 设置最多返回的关键点数量。
                      # qualityLevel : 角点的品质因子，决定能不能成为角点。
                      # minDistance : 关键点之间的最少像素点，此邻域范围内如果存在更强角点，则删除此角点。
                      # blockSize : 计算一个角点是否为关键点时所取的区域大小。
# lucas kanade光流法参数设定
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
                 # winSize：每个金字塔级别的搜索窗口大小。
                 # maxLevel：从0开始的最大金字塔等级数字；如果设置为0，则不使用金字塔（单个级别）；如果设置为1，则使用两个级别。
                 # criteria：参数，指定迭代搜索算法的终止条件（在指定的最大迭代次数之后 criteria.maxCount 或当搜索窗口移动小于时 criteria.epsilon）。

cap = cv2.VideoCapture(videoFileName)

#计算第一帧特征点
ret, prev = cap.read()
prevGray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
p0 = cv2.goodFeaturesToTrack(prevGray, mask=None, **feature_params)

while True:
    ret, frame = cap.read()
    if not  ret: #没读到当前帧
        break

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    #计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)

    #选取好的跟踪点
    goodPoints = p1[st==1]
    goodPrevPoints = p0[st==1]
    # 在结果图像中迭加画出特征点和计算出来的光流向量
    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv2.line(res, (x0, y0), (x1, y1), drawColor)
        cv2.circle(res, (x0, y0), 3, drawColor)

    #更新上一帧
    prevGray = gray.copy()
    p0 = goodPoints.reshape(-1,1,2)

    #显示计算结果图像
    cv2.imshow('计算结果', res)

    key = cv2.waitKey(30)  #每一帧间隔30ms
    if key==27:
        break

###################################   代码二    ##########################################

import cv2
import numpy as np
# 加载视频
cap = cv2.VideoCapture()
cap.open(r'E:\deep learning\cv\data\vtest.avi')
if not cap.isOpened():
    print("无法打开视频文件")

# 方便调参，使用dict字典
# lucas kanade光流法参数设定
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
                 # winSize：每个金字塔级别的搜索窗口大小。
                 # maxLevel：从0开始的最大金字塔等级数字；如果设置为0，则不使用金字塔（单个级别）；如果设置为1，则使用两个级别。
                 # criteria：参数，指定迭代搜索算法的终止条件（在指定的最大迭代次数之后 criteria.maxCount 或当搜索窗口移动小于时 criteria.epsilon）。
# 角点检测参数设定
feature_params = dict(maxCorners=500,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)
                      # maxCorners : 设置最多返回的关键点数量。
                      # qualityLevel : 角点的品质因子，决定能不能成为角点。
                      # minDistance : 关键点之间的最少像素点，此邻域范围内如果存在更强角点，则删除此角点。
                      # blockSize : 计算一个角点是否为关键点时所取的区域大小。

class App:
    def __init__(self, video_src):  # 构造方法，初始化一些参数和视频路径
        self.track_len = 10
        self.detect_interval = 5
        self.tracks = []
        self.cam = cv2.VideoCapture(video_src)
        self.frame_idx = 0

    def run(self):  # 光流运行方法
        while True:
            ret, frame = self.cam.read()  # 读取视频帧
            if ret:
                frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # 转化为灰度虚图像
                vis = frame.copy()

                if len(self.tracks) > 0:  # 检测到角点后进行光流跟踪
                    img0, img1 = self.prev_gray, frame_gray
                    p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2)
                    p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None,
                                                           **lk_params)  # 前一帧的角点和当前帧的图像作为输入来得到角点在当前帧的位置
                    p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None,
                                                            **lk_params)  # 当前帧跟踪到的角点及图像和前一帧的图像作为输入来找到前一帧的角点位置
                                                                          # **用法，传参，可以将定义好的参数直接使用
                    """
                    nextPts, status, err = calcOpticalFlowPyrLK(prevImg, nextImg, prevPts[, nextPts[, status[,
                    err[, winSize[, maxLevel[, criteria[, flags[, minEigThreshold]]]]]]]])
                    参数说明：
                      prevImage 前一帧8-bit图像
                      nextImage 当前帧8-bit图像
                      prevPts 待跟踪的特征点向量
                      nextPts 输出跟踪特征点向量
                      status 特征点是否找到，找到的状态为1，未找到的状态为0
                      err 输出错误向量，（不太理解用途...）
                      winSize 搜索窗口的大小
                      maxLevel 最大的金字塔层数
                      flags 可选标识：OPTFLOW_USE_INITIAL_FLOW   OPTFLOW_LK_GET_MIN_EIGENVALS
                    """
                    d = abs(p0 - p0r).reshape(-1, 2).max(-1)  # 得到角点回溯与前一帧实际角点的位置变化关系
                    good = d < 1  # 判断d内的值是否小于1，大于1跟踪被认为是错误的跟踪点
                    new_tracks = []
                    for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good):  # 将跟踪正确的点列入成功跟踪点
                        if not good_flag:
                            continue
                        tr.append((x, y))
                        if len(tr) > self.track_len:
                            del tr[0]
                        new_tracks.append(tr)
                        cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)
                    self.tracks = new_tracks
                    cv2.polylines(vis, [np.int32(tr) for tr in self.tracks], False,
                                  (0, 255, 0))  # 以上一振角点为初始点，当前帧跟踪到的点为终点划线
                    # draw_str(vis, (20, 20), 'track count: %d' % len(self.tracks))

                if self.frame_idx % self.detect_interval == 0:  # 每5帧检测一次特征点
                    mask = np.zeros_like(frame_gray)  # 初始化和视频大小相同的图像
                    mask[:] = 255  # 将mask赋值255也就是算全部图像的角点
                    for x, y in [np.int32(tr[-1]) for tr in self.tracks]:  # 跟踪的角点画圆
                        cv2.circle(mask, (x, y), 5, 0, -1)
                    p = cv2.goodFeaturesToTrack(frame_gray, mask=mask, **feature_params)  # 像素级别角点检测
                    if p is not None:
                        for x, y in np.float32(p).reshape(-1, 2):
                            self.tracks.append([(x, y)])  # 将检测到的角点放在待跟踪序列中

                self.frame_idx += 1
                self.prev_gray = frame_gray
                cv2.imshow('lk_track', vis)

            ch = 0xFF & cv2.waitKey(100)
            if ch == 27:
                break

App(r'E:\deep learning\cv\data\vtest.avi').run()
